from typing import TypedDict, Annotated, List, Dict, Any, Optional
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from langgraph.types import interrupt, Command
from langgraph.checkpoint.memory import MemorySaver
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage, AnyMessage
from langchain_openai import ChatOpenAI
import os
from datetime import datetime
import json

# 初始化大模型
model = ChatOpenAI(
    model_name="qwen3",
    openai_api_key="sk-c1",
    openai_api_base="http://172.16.53.70:31963/v1",
)

# 订票系统的状态定义
class BookingState(TypedDict):
    messages: Annotated[List[AnyMessage], add_messages]  # LangGraph 消息模式
    user_query: str
    destination: str
    travel_date: str
    passenger_count: int
    ticket_class: str
    price: float
    booking_confirmed: bool
    current_step: str
    user_choice: Optional[str]  # 中断响应的用户选择


# 节点1: 订票信息收集
async def collect_booking_info(state: BookingState) -> Dict[str, Any]:
    """收集订票所需的基本信息"""
    print("📝 正在收集订票信息...")
    
    # 解析用户查询，提取订票信息
    user_query = state["user_query"]
    
    # 使用大模型提取订票信息
    system_prompt = """
    你是一个机票预订助手，需要从用户的查询中提取以下信息：
    1. 目的地
    2. 出行日期
    3. 乘客人数
    4. 舱位等级（经济舱、商务舱、头等舱）
    
    请以JSON格式返回提取的信息，格式如下：
    {
        "destination": "目的地城市名",
        "travel_date": "YYYY-MM-DD格式的日期",
        "passenger_count": 乘客数量(整数),
        "ticket_class": "舱位等级"
    }
    
    如果某项信息未提供，请使用合理的默认值。
    """
    
    messages = [
        SystemMessage(content=system_prompt),
        HumanMessage(content=user_query)
    ]
    
    # 调用大模型提取信息
    response = await model.ainvoke(messages)
    
    try:
        # 尝试解析JSON响应
        # 首先尝试从响应中提取JSON部分
        content = response.content
        # 查找可能的JSON内容（通常在```json和```之间，或者整个内容就是JSON）
        json_content = content
        
        # 尝试查找markdown格式的JSON代码块
        import re
        json_match = re.search(r'```(?:json)?\s*([\s\S]*?)\s*```', content)
        if json_match:
            json_content = json_match.group(1)
        
        # 清理JSON字符串，移除可能的非JSON字符
        json_content = json_content.strip()
        
        # 解析JSON
        extracted_info = json.loads(json_content)
        destination = extracted_info.get("destination", "北京")
        travel_date = extracted_info.get("travel_date", "2023-12-25")
        passenger_count = extracted_info.get("passenger_count", 1)
        ticket_class = extracted_info.get("ticket_class", "经济舱")
    except Exception as e:
        print(f"解析大模型响应时出错: {str(e)}")
        print(f"原始响应内容: {response.content}")
        # 使用默认值
        destination = "北京"
        travel_date = "2023-12-25"
        passenger_count = 1
        ticket_class = "经济舱"
    
    # 根据目的地和舱位等级计算票价
    # 这里使用简单的逻辑，实际应用中可能需要查询数据库或调用外部API
    base_prices = {
        "北京": 1000,
        "上海": 800,
        "广州": 1200,
        "深圳": 1100,
        "成都": 1500
    }
    
    class_multipliers = {
        "经济舱": 1.0,
        "商务舱": 2.5,
        "头等舱": 4.0
    }
    
    base_price = base_prices.get(destination, 1000)
    multiplier = class_multipliers.get(ticket_class, 1.0)
    price = base_price * multiplier
    
    # 将当前用户问题加入消息
    new_messages = [HumanMessage(content=state["user_query"])]
    
    return {
        **state,  # 保留所有现有状态
        "messages": new_messages,
        "destination": destination,
        "travel_date": travel_date,
        "passenger_count": passenger_count,
        "ticket_class": ticket_class,
        "price": price,
        "current_step": "booking_confirmation",
        "booking_confirmed": False
    }


# 节点2: 订票确认中断
async def booking_confirmation_interrupt(state: BookingState) -> Dict[str, Any]:
    """中断以获取用户对订票信息的确认"""
    print("🔍 正在确认订票信息...")
    
    # 构建确认信息
    confirmation_msg = f"""## 订票确认
    
我已收到您的订票请求，请确认以下信息：

- 目的地: {state['destination']}
- 出行日期: {state['travel_date']}
- 乘客人数: {state['passenger_count']}人
- 舱位等级: {state['ticket_class']}
- 票价: ¥{state['price']}

请选择：
- **confirm**: 确认订票
- **cancel**: 取消订票

您是否确认上述订票信息？"""

    user_choice = interrupt(confirmation_msg)
    
    print(f"DEBUG: 用户在 booking_confirmation_interrupt 中选择: {user_choice}")
    
    return {
        **state,  # 保留所有现有状态
        "user_choice": user_choice,
        "current_step": "process_booking",
    }


# 节点3: 处理订票
async def process_booking(state: BookingState) -> Dict[str, Any]:
    """根据用户确认处理订票"""
    print("🎫 正在处理订票...")
    
    user_choice = state.get("user_choice", "cancel")
    booking_confirmed = user_choice == "confirm"
    
    # 使用大模型生成更自然的回复
    system_prompt = """
    你是一个专业的机票预订助手。请根据用户的选择生成适当的回复。
    如果用户确认订票，请生成一个友好的确认信息，包括订票详情。
    如果用户取消订票，请生成一个礼貌的取消确认信息。
    回复应该简洁、专业且友好。
    """
    
    booking_info = {
        "confirmed": booking_confirmed,
        "destination": state["destination"],
        "travel_date": state["travel_date"],
        "passenger_count": state["passenger_count"],
        "ticket_class": state["ticket_class"],
        "price": state["price"]
    }
    
    user_message = f"用户{'确认' if booking_confirmed else '取消'}了以下订票：目的地 {state['destination']}，出行日期 {state['travel_date']}，{state['passenger_count']}人，{state['ticket_class']}，票价 ¥{state['price']}。"
    
    messages = [
        SystemMessage(content=system_prompt),
        HumanMessage(content=user_message)
    ]
    
    # 调用大模型生成回复
    try:
        response = await model.ainvoke(messages)
        assistant_message = response.content
    except Exception as e:
        print(f"生成回复时出错: {str(e)}")
        # 使用备用回复
        if booking_confirmed:
            assistant_message = f"""
订票成功！您的订单已确认。

以下是您的订票信息：
- 目的地: {state["destination"]}
- 出行日期: {state["travel_date"]}
- 乘客人数: {state["passenger_count"]}人
- 舱位等级: {state["ticket_class"]}
- 票价: ¥{state["price"]}

感谢您使用我们的订票系统！
"""
        else:
            assistant_message = """
订票已取消。感谢您的访问。

期待您的再次光临！
"""
    
    # 将助手回复加入消息
    new_messages = [AIMessage(content=assistant_message)]
    
    return {
        **state,  # 保留所有现有状态
        "messages": new_messages,
        "booking_confirmed": booking_confirmed,
        "current_step": "completed",
    }


# 构建订票系统流程图
async def create_graph(checkpointer):
    """创建订票系统工作流"""
    builder = StateGraph(BookingState)
    
    # 添加节点
    builder.add_node("collect_booking_info", collect_booking_info)
    builder.add_node("booking_confirmation_interrupt", booking_confirmation_interrupt)
    builder.add_node("process_booking", process_booking)
    
    # 添加边
    builder.add_edge(START, "collect_booking_info")
    builder.add_edge("collect_booking_info", "booking_confirmation_interrupt")
    builder.add_edge("booking_confirmation_interrupt", "process_booking")
    builder.add_edge("process_booking", END)
    
    return builder.compile(checkpointer=checkpointer)



# 订票系统流式响应函数
async def stream_booking_response(checkpointer, thread_id: str, user_choice: str):
    """流式输出订票流程"""
    print(f"\n📝 开始流式订票，线程: {thread_id}")
    print(f"用户选择: {user_choice}")
    
    config = {"configurable": {"thread_id": thread_id}}
    
    try:
        # 确保图对象已初始化
        graph = await create_graph(checkpointer=checkpointer)
        
        # 用 Command 恢复流程，astream_events 捕获流式内容
        events_received = False
        async for event in graph.astream_events(
            Command(resume=user_choice), version="v2", config=config
        ):
            events_received = True
            # 处理流式事件
            if event["event"] == "on_chat_model_stream":
                data = event["data"]
                chunk = data["chunk"]
                if hasattr(chunk, "content") and chunk.content:
                    yield {
                        "content": chunk.content,
                        "type": "content",
                        "done": False,
                    }
        
        # 如果没有收到任何事件，使用大模型生成流式响应
        if not events_received:
            # 获取会话状态
            try:
                # 尝试获取会话状态
                thread_config = {"configurable": {"thread_id": thread_id}}
                state = await checkpointer.get(thread_id)
                
                # 如果成功获取状态，使用状态中的信息
                if state and isinstance(state, dict):
                    destination = state.get("destination", "北京")
                    travel_date = state.get("travel_date", "2023-12-25")
                    passenger_count = state.get("passenger_count", 1)
                    ticket_class = state.get("ticket_class", "经济舱")
                    price = state.get("price", 1200.0)
                else:
                    # 使用默认值
                    destination = "北京"
                    travel_date = "2023-12-25"
                    passenger_count = 1
                    ticket_class = "经济舱"
                    price = 1200.0
            except Exception as e:
                print(f"获取会话状态时出错: {str(e)}")
                # 使用默认值
                destination = "北京"
                travel_date = "2023-12-25"
                passenger_count = 1
                ticket_class = "经济舱"
                price = 1200.0
            
            # 使用大模型生成响应
            system_prompt = """
            你是一个专业的机票预订助手。请根据用户的选择生成适当的回复。
            如果用户确认订票，请生成一个友好的确认信息，包括订票详情。
            如果用户取消订票，请生成一个礼貌的取消确认信息。
            回复应该简洁、专业且友好。
            """
            
            user_message = f"用户{'确认' if user_choice == 'confirm' else '取消'}了以下订票：目的地 {destination}，出行日期 {travel_date}，{passenger_count}人，{ticket_class}，票价 ¥{price}。"
            
            messages = [
                SystemMessage(content=system_prompt),
                HumanMessage(content=user_message)
            ]
            
            # 使用流式模式调用大模型
            try:
                async for chunk in model.astream(messages):
                    if hasattr(chunk, "content") and chunk.content:
                        yield {
                            "content": chunk.content,
                            "type": "content",
                            "done": False,
                        }
            except Exception as e:
                print(f"流式生成回复时出错: {str(e)}")
                # 使用备用回复
                if user_choice == "confirm":
                    booking_result = f"""
订票成功！您的订单已确认。

以下是您的订票信息：
- 目的地: {destination}
- 出行日期: {travel_date}
- 乘客人数: {passenger_count}人
- 舱位等级: {ticket_class}
- 票价: ¥{price}

感谢您使用我们的订票系统！
"""
                else:
                    booking_result = """
订票已取消。感谢您的访问。

期待您的再次光临！
"""
                
                yield {
                    "content": booking_result,
                    "type": "content",
                    "done": False,
                }
        
        # 最终完成信号
        yield {"content": "", "type": "content", "done": True}
        
    except Exception as e:
        print(f"流式错误: {str(e)}")
        yield {
            "content": f"流式错误: {str(e)}",
            "type": "error",
            "done": True,
        }


# 启动订票流程
async def start_booking_process(checkpointer, user_query: str, thread_id: str = None):
    """启动订票流程"""
    if thread_id is None:
        thread_id = f"booking_{datetime.now().strftime('%Y%m%d%H%M%S')}"
    
    # 初始化状态
    initial_state = {
        "messages": [],
        "user_query": user_query,
        "destination": "",
        "travel_date": "",
        "passenger_count": 0,
        "ticket_class": "",
        "price": 0.0,
        "booking_confirmed": False,
        "current_step": "",
        "user_choice": None
    }
    
    # 配置线程ID
    config = {"configurable": {"thread_id": thread_id}}
    
    # 确保图对象已初始化
    graph = await create_graph(checkpointer=checkpointer)
    
    # 启动工作流
    await graph.ainvoke(initial_state, config=config)
    
    return thread_id
